# coding: utf-8

# Copyright (c) 2020 CSTCHC Technologies Co., Ltd.
# cstc-hardware-cap is licensed under the Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#     http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
# PURPOSE.
# See the Mulan PSL v2 for more details.

"""card test"""
import os
import subprocess
import argparse
import time
import commands

from cstchccompatible.command import Command, CertCommandError
from cstchccompatible.test import Test

class Card_test(Test):
    """
       modules test
    """
    def __init__(self):
        Test.__init__(self)
        self.args = None
        self.device = None
        self.requirements = ["gcc-c++"]

    def setup(self, args=None):
        """
        Initialization before test
        :param args:
        :return:
        """
        self.args = args or argparse.Namespace()
        self.device = getattr(args, 'device', None)
   
    def current_card(self):
        print("Vendor Info:")
        pci_num = self.device.get_property("DEVPATH").split('/')[-1]
        Command('lspci -s %s -v' % pci_num).echo()
        
        print("Drive Info:")
        Command('modinfo nvidia|head -n 13').echo()
        return pci_num

    def modules(self):
        if commands.getoutput('lsmod |grep nouveau') != '':
            try:
                Command('rmmod nouveau').echo()
            except:
                return False

    def pressure(self):
        pci_key = self.current_card()
        print("Driver is installed\n")

        print("Monitor GPU temperature and utilization rate")
        gpu_connect = Command('nvidia-smi -q')
        gpu_connect._run()
        
        output = {}
        dicts = {}
        gpu = gpu_connect.output
        gpu.append('')
        for line in gpu:
            if line.strip() == '':
                for pcis in dicts.values():
                    if pci_key in pcis:
                        output[pci_key] = {}
                        output[pci_key].update(dicts)
            else:
                if ':' in line:
                    key = line.split(':', 1)[0].strip()
                    value = line.split(':', 1)[1].strip()
                    dicts[key] = value
        
        id_num = output[pci_key]['Minor Number']
        os.environ['CUDA_VISIBLE_DEVICES'] = id_num
        
        
        os.system('cd /usr/bin && nohup gpu_burn 10 &')
        # Pressure measurement
        while not subprocess.call("ps -ef |grep 'gpu_burn' | grep -v grep", shell=True):
            Command('nvidia-smi').echo()
            time.sleep(1)
        
    def samples(self):
        print('Run test case')

        if commands.getoutput('arch') == 'aarch64':
            os.chdir('/var/cuda_samples/NVIDIA_CUDA-11.4_Samples/bin/sbsa/linux/release')
        else:
            os.chdir('/var/cuda_samples/NVIDIA_CUDA-11.4_Samples/bin/x86_64/linux/release')

        sample_list = ['simpleOccupancy', 'bandwidthTest', 'p2pBandwidthLatencyTest', 'deviceQuery', 'clock']
        for sample in sample_list:
            for files in os.popen('ls'):
                line = files.strip()
                if line == sample:
                    print('INFO: ==================     %s     ===============' % line)
                    message = os.popen('./%s' % line).readlines()
                    for mes in message:
                        if 'PASS' in mes:
                            print(' %s is passed' % line)
    
    def test(self):
        try:
            self.modules()
            self.pressure()
            # Check whether the pressure test is successful
            if os.path.exists('/usr/bin/nohup.out'):
                for content in open('/usr/bin/nohup.out'):
                   line = content.strip()
                   if "FAULTY" in line:
                       print("Pressure test failed\n")
                       return False
            self.samples()
            return True
        except Exception as e:
            print(e)
            return False
